#!/usr/bin/perl -w

###############################################################################
# Export a ped file and a map file
#
# Parameters needed: dbname, username, temp_dir, ped_file, map_file, 
#                    idlink_filter, pheno_filter, anno_filter
#
# Author: Herodotos Herodotou and Nedyalko Borisov
# Date: May 4, 2010
#
##############################################################################

chdir("./src"); # Change the execution directory to src
use lib './src'; # Specify the library location
use getOpts qw(get_options println);
sub remove_file;
sub remove_tmp_tables;

if ($#ARGV != 0)
{
   println qq(Please specify the .properties file);
   println qq(Usage: perl $0 export.properties);
   exit(-1);
}

# Get input parameters
%opt = get_options($ARGV[0]);
$DB_NAME  = $opt{'dbname'};
$USERNAME = $opt{'username'};
$TEMP_DIR = $opt{'temp_dir'};
$PED_FILE = $opt{'ped_file'};
$MAP_FILE = $opt{'map_file'};

$IDLINK_FILTER = $opt{'idlink_filter'};
$PHENO_FILTER  = $opt{'pheno_filter'};
$ANNO_FILTER   = $opt{'anno_filter'};


# Create the temp directory
if (! -d "$TEMP_DIR")
{
   mkdir $TEMP_DIR or die $!;
}

# Ask for confirmation and remove existing files
remove_file($PED_FILE);
remove_file($MAP_FILE);

$start = time();
println qq(Starting export preprocessing);

# Remove the temporary tables if they exist
remove_tmp_tables();

#####################################################################
# Filtering
#####################################################################

# Build a subquery involving the pheno and idlink filters
$subquery = qq(FROM idlink i JOIN study_platform sp ON lower\(sp.study_id\) = lower\(i.study_id\) );

if ($PHENO_FILTER)
{
   # The user provided filters on pheno table
   $subquery = qq($subquery JOIN pheno p ON lower\(p.patient_id\) = lower\(i.patient_id\) );
  
   if ($PHENO_FILTER =~ /feature/ || $PHENO_FILTER =~ /value/)
   {
      # The user provided filters on pheno_details table
      $subquery = qq($subquery JOIN pheno_details pd ON lower\(p.patient_id\) = lower\(pd.patient_id\) );
   }

   $PHENO_FILTER =~ s/patient_id/p\.patient_id/gi;
   $subquery = qq($subquery WHERE \( $PHENO_FILTER \) );
}

if ($IDLINK_FILTER)
{
   # The user provided filters on idlink table
   $keyword = ($PHENO_FILTER) ? "AND" : "WHERE";
   $IDLINK_FILTER =~ s/patient_id/i\.patient_id/gi;
   $IDLINK_FILTER =~ s/study_id/i\.study_id/gi;
   $subquery = qq($subquery $keyword \( $IDLINK_FILTER \) );
}

# Gather all idlink_ids and corresponding patient_ids and platforms
$sql = qq(SELECT i.idlink_id, i.patient_id, i.exp_id, sp.platform $subquery);
!system qq(psql -U $USERNAME -d $DB_NAME -t -A -F "|" ). 
        qq( -c "$sql" -o $TEMP_DIR/idlinks.txt ) or die "ERROR: Invalid query:\n$sql";

# Check if we have any patient data to export
if (-z "$TEMP_DIR/idlinks.txt")
{
   println qq(No patients matching your filter conditions were found);
   !system qq(rm -f $TEMP_DIR/idlinks.txt) or die $!;
   exit(0);
}

#####################################################################
# Merging across studies
#####################################################################

# Gather all exported study_ids into a temp table
$sql = qq(CREATE TABLE tmp_studies \( study_id VARCHAR\(20\) \));
!system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 ) or die $!;

$sql = qq(INSERT INTO tmp_studies\( study_id \)
          SELECT DISTINCT lower\(i.study_id\) $subquery);
!system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null ) or die $!;

# Determine the snps that are common for all exporting studies
$study_count_a = 0;
$study_count_i = 0;

# Create filter table for anno_ids based on the selected studies
$sql = qq(CREATE TABLE tmp_study_anno_ids\( anno_id INTEGER, platform CHAR\(1\)\) );
!system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 ) or die $!;

# Get one idlink from each study
$sql = qq( SELECT  MIN\(idlink_id\), platform
           FROM    tmp_studies s
           JOIN    study_platform sp ON lower\(s.study_id\) = lower\(sp.study_id\)
           JOIN    idlink i ON lower\(i.study_id\) = lower\(s.study_id\)
           GROUP BY s.study_id, platform 
       );

!system qq(psql -U $USERNAME -d $DB_NAME -t -A -F "|" -c "$sql" ) .
        qq( -o $TEMP_DIR/filter_idlinks > /dev/null 2>&1 ) or die $!;

# Loop through all the idlinks from different studies and merge their 
# corresponding anno_ids
open INFILE, "<", $TEMP_DIR."/filter_idlinks" or die $!;
while ($line = <INFILE>)
{
   # Skip empty lines
   if ($line =~ /^\s*$/)
   {
      next;
   }

   # Get the idlink_id and platform 
   @data = split(/\|/, $line);
   $idlink_id = $data[0];
   $platform = $data[1];

   $idlink_id =~ s/\s+//g;
   $platform =~ s/\s+//g;

   $sql = qq(INSERT INTO tmp_study_anno_ids\( anno_id, platform \)
             SELECT anno_id, '$platform' 
             FROM   geno_$idlink_id
           );
   !system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 ) or die $!;
   if( $platform eq 'i' )
   {
     $study_count_i++;
   }
   else
   {
     $study_count_a++;
   }
}

close INFILE;
system qq(rm -f $TEMP_DIR/filter_idlinks);


#####################################################################
# Merging across platforms
#####################################################################

# Create the anno filter tables
$sql = qq( CREATE TABLE tmp_filter_anno_i\( anno_id INTEGER, 
              alleleA CHAR\(1\), alleleB CHAR\(1\) \);
           CREATE TABLE tmp_filter_anno_a\( anno_id INTEGER, 
              alleleA CHAR\(1\), alleleB CHAR\(1\) \);
          );
!system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 ) or die $!;

# Prepare the anno filter
$and_anno_filter = "";
if ($ANNO_FILTER)
{
   $and_anno_filter = qq( AND \( $ANNO_FILTER \) );
}

# Determine if platform merging is required
if( $study_count_i > 0 && $study_count_a > 0 )
{
  # Merge platforms
  for $plat ("i", "a")
  {
    $sql = qq( INSERT INTO tmp_filter_anno_$plat \(anno_id, alleleA, alleleB\) 
               SELECT  unp.anno_id, alleleA, alleleB 
               FROM    \(
                          SELECT MIN\(a.anno_id\) AS anno_id, a.rsid, a.alleleA, a.alleleB 
                          FROM   anno a
                          WHERE  platform = '$plat' $and_anno_filter
                          GROUP  BY rsid, platform, alleleA, alleleB 
                      \) unp 
                JOIN \(  
                        SELECT rsid
                        FROM   anno a
                        JOIN \(
                                SELECT anno_id
                                FROM   tmp_study_anno_ids 
                                WHERE  platform = 'i'
                                GROUP  BY anno_id 
                                HAVING COUNT\(anno_id\) = $study_count_i
                                UNION ALL
                                SELECT anno_id
                                FROM   tmp_study_anno_ids 
                                WHERE  platform = 'a'
                                GROUP  BY anno_id 
                                HAVING COUNT\(anno_id\) = $study_count_a
                              \) sa ON sa.anno_id = a.anno_id 
                        GROUP BY rsid 
                        HAVING SUM\( CASE WHEN a.platform = 'i' THEN 1 ELSE 1000 END\) > 1000 
                           AND SUM\( CASE WHEN a.platform = 'i' THEN 1 ELSE 1000 END \)% 1000 > 0 
                    \) unr ON unp.rsid = unr.rsid 
              );

    !system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null) or die;
  }
}
else
{
   # No merge, only one platform
   if ($study_count_i > 0) {
      $platform = "i";
      $study_count = $study_count_i;
   }
   else
   {
      $platform = "a";
      $study_count = $study_count_a;
   }

   $sql = qq( INSERT INTO tmp_filter_anno_$platform \(anno_id, alleleA, alleleB\) 
              SELECT a.anno_id, a.alleleA, a.alleleB 
              FROM anno a 
              JOIN \(
                     SELECT anno_id
                    FROM   tmp_study_anno_ids 
                    WHERE  platform = '$platform'
                    GROUP  BY anno_id 
                    HAVING COUNT\(anno_id\) = $study_count 
                  \) sa ON sa.anno_id = a.anno_id
             WHERE platform = '$platform' $and_anno_filter);

  !system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null) or die;
}

# Check if we have any snp data to export
$have_snp_data = 0;
for $plat ("i", "a")
{
   $sql = qq(SELECT anno_id FROM tmp_filter_anno_$plat LIMIT 1);
   !system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" -t -A).
           qq( -o $TEMP_DIR/anno_data_$plat.txt > /dev/null 2>&1 ) or die $!;
   $have_snp_data = $have_snp_data + 1 if not (-z "$TEMP_DIR/anno_data_$plat.txt");
   !system qq(rm -f $TEMP_DIR/anno_data_$plat.txt) or die $!;
}
if ($have_snp_data == 0)
{
   println qq(No SNPs matching your filter conditions were found);
   exit(0);
}

$pre_time = time() - $start;
println qq(Completed export preprocessing);

#####################################################################
# PED exporting
#####################################################################

println qq(Starting PED export);
$start = time();

# Perform the PED export
$num_patients = 0;
open INFILE, "<", $TEMP_DIR."/idlinks.txt" or die $!;

# Create the family id sequence
$sql = qq(DROP SEQUENCE family_id);
system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 );
$sql = qq(CREATE SEQUENCE family_id START 1);
!system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 ) or die $!;

while ($line = <INFILE>)
{
   # Skip empty lines
   if ($line =~ /^\s*$/)
   {
      next;
   }

   # Get the idlink_id and patiend_id pair
   @data = split(/\|/, $line);
   $idlink_id = $data[0];
   $patient_id = $data[1];
   $exp_id = $data[2];
   $platform = $data[3];

   $idlink_id =~ s/\s+//g;
   $patient_id =~ s/\s+//g;
   $platform =~ s/\s+//g;

   # Execute the query
   $query = qq(
      SELECT    nextval('family_id') || E'\t' || '$exp_id'  || E'\t0\t0\t' || 
                  CASE 
                    WHEN UPPER\(p.sex\) = 'M' THEN 1 
                    WHEN UPPER\(p.sex\) = 'F' THEN 2 
                    ELSE 0 
                  END
                   ||  E'\t0\t' as prefix,
                   array_to_string\( array\( 
                        SELECT
                            CASE 
                                WHEN snpval = 0 THEN allelea || allelea
                                WHEN snpval = 1 THEN allelea || alleleb
                                WHEN snpval = 2 THEN alleleb || alleleb
                                WHEN snpval = -9 THEN '--'
                                ELSE '--'
                            END
                        FROM  geno_$idlink_id g 
                        JOIN  tmp_filter_anno_$platform fa ON fa.anno_id = g.anno_id
                        ORDER BY fa.anno_id 
                        \), E'\t' 
                \)
        FROM     pheno p 
        WHERE    p.patient_id = '$patient_id';);

   !system qq(psql -d $DB_NAME -U $USERNAME -t).
           qq( -A -F "" -c "$query" >> $PED_FILE) or die $!;
   ++$num_patients;
}

# Drop the family id sequence
$sql = qq(DROP SEQUENCE family_id);
!system qq(psql -U $USERNAME -d $DB_NAME -c "$sql" > /dev/null 2>&1 ) or die $!;

# Done with PED export
close INFILE;
system qq(rm -f $TEMP_DIR/idlinks.txt);

#####################################################################
# MAP exporting
#####################################################################

$ped_time = time() - $start;
println qq(Export to $PED_FILE completed);

println qq(Starting MAP export);
$start = time();


# Export the MAP file
$platform = ($study_count_i > 0) ? "i" : "a";
$query = qq(
            SELECT chromosome || E'\t' || rsid || E'\t' || 0 || E'\t' || location
            FROM   anno a
            JOIN   tmp_filter_anno_$platform fa ON fa.anno_id = a.anno_id
            ORDER BY fa.anno_id
           );

!system qq(psql -d $DB_NAME -U $USERNAME -t).
        qq( -A -F "" -c "$query" -o $MAP_FILE > /dev/null) or die $!;

$map_time = time() - $start;
println qq(Export to $PED_FILE completed);

#####################################################################
# Clean up and reports
#####################################################################

# Remove the temporary tables
remove_tmp_tables();

# Remove the temp directory if empty
if (! scalar <$TEMP_DIR/*>)
{
   rmdir $TEMP_DIR;
}

# Print stats
println qq();
$total_time = $pre_time + $ped_time + $map_time;
println qq(Number of patients:\t $num_patients);
println qq(Time for preprocessing (sec):\t $pre_time);
println qq(Time for PED export (sec):\t $ped_time);
println qq(Time for MAP export (sec):\t $map_time);
println qq(Total Time to export (sec):\t $total_time);


###############################################################################
# Function: remove_file
#
# Input: filePath 
#
# Description: if file exists, confirms file removal or exits
#
##############################################################################
sub remove_file()
{
  my $FILE = $_[0];
  if (-e $FILE)
  {
    println "$FILE already exists. This process will replace it.";
    print "Are you sure (yes/no): ";

    $| = 1;
    $_ = <STDIN>;
    chomp;

    if ( $_ ne "yes" )
    {
      exit(0);
    }

    # Remove the existing file
    !system qq(rm -f $FILE) or die $!;
  }
}


###############################################################################
# Function: remove_tmp_tables
#
# Input: none 
#
# Description: Removes all temporary tables thar are used by the export script
#
##############################################################################
sub remove_tmp_tables()
{
   system qq(psql -U $USERNAME -d $DB_NAME).
          qq( -c "DROP TABLE IF EXISTS tmp_studies" > /dev/null 2>&1 );
  system qq(psql -U $USERNAME -d $DB_NAME).
          qq( -c "DROP TABLE IF EXISTS tmp_study_anno_ids" > /dev/null 2>&1 );
  system qq(psql -U $USERNAME -d $DB_NAME).
          qq( -c "DROP TABLE IF EXISTS tmp_filter_anno_i" > /dev/null 2>&1 );
  system qq(psql -U $USERNAME -d $DB_NAME).
          qq( -c "DROP TABLE IF EXISTS tmp_filter_anno_a" > /dev/null 2>&1 );
}
